# -*- coding: utf-8 -*-
"""
Created on Fri Sep 01 10:41:11 2017

@author: za-caoyixin
"""
import numpy as np
import pandas as pd


def strategy(mktdata, sgn_buy, sgn_sell, sgn_coverbuy, sgn_coversell, buymode = True, sellmode = False, stoplos = False, stopwin = False, logmode=False,stoplospoint=0.0,**kw):
    #%%
    if 'stoplospoint' in kw:
        stoplospoint = kw['stoplospoint']
    if 'stoplospoint' not in kw:
        stoplospoint=np.zeros(len(mktdata))
    if 'stopwin' in kw:
        stopwin = kw['stopwin']
    if 'stoplos' in kw:
        stoplos = kw['stoplos']
    if 'sellmode' in kw:
        sellmode = kw['sellmode']
    if 'buymode' in kw:
        buymode = kw['buymode']
    
    """模拟交易过程"""
    sgn_en = np.zeros(len(mktdata))     # entry signals, trade direciton: 1: buy, -1: sell, 0, do nothing
    sgn_ex = np.zeros(len(mktdata))     # exit signals, cover direction: 1: coversell, -1: coverbuy
    sgn_hd = np.zeros(len(mktdata))     # hold position: 1: long, -1:short, 0, no position
    prc_en = np.zeros(len(mktdata))     # price of opening a position
    prc_ex = np.zeros(len(mktdata))     # price of covering a position

    popen = mktdata['open'].values
    pclose = mktdata['close'].values
    phigh = mktdata['high'].values
    plow = mktdata['low'].values
    
    idx_en = np.array([],dtype=int)     # 开仓点
    idx_ex = np.array([],dtype=int)     # 平仓点
    #%%
    
    
    for ii in np.arange(len(mktdata)):
        #%%
        if ii > 0:
            
            # 执行前一根K线发出的开仓信号
            if (sgn_en[ii-1] == 1) or (sgn_en[ii-1] == -1):
                sgn_hd[ii] = sgn_en[ii-1]                           # 更新当前k线持仓状态
                prc_en[ii] = - np.sign(sgn_hd[ii]) * pclose[ii-1]   # 将前一根K线收盘价作为开仓价，若多头，开仓价设为负，空头，设为正
                prc_ex[ii] =   np.sign(sgn_hd[ii]) * pclose[ii]      # 当前K线收盘价作为浮动平仓价格
                idx_en = np.append(idx_en, [ii])
                highAfterEntry = phigh[ii]
                lowAfterEntry = plow[ii]
#                longstop = min([plow[ii], plow[ii-1]]) - stoplospoint   #设置止损多头价位为最近两根k线最低价减去一定点数
#                shortstop = max([phigh[ii], phigh[ii-1]]) + stoplospoint 
                longstop = min([plow[ii], plow[ii-1]]) - stoplospoint[ii-1]
                shortstop = max([phigh[ii], phigh[ii-1]]) + stoplospoint[ii-1]
                if logmode:
                    print('开仓,方向%s, 开仓价%s, longstop:%s, shortstop:%s' % (sgn_en[ii-1], prc_en[ii], longstop, shortstop))
            # 如果当下k线已持仓，判断当下k线是否在收盘发出平仓信号
            if (sgn_hd[ii] == 1) and sgn_coverbuy[ii]:
                sgn_ex[ii] = -1                                 # 平多仓
                idx_ex = np.append(idx_ex, [ii])
                if logmode:
                    print('平多仓, 平仓价%s, longstop:%s, shortstop:%s' % (prc_ex[ii], longstop, shortstop))
            elif (sgn_hd[ii] == -1) and sgn_coversell[ii]:
                sgn_ex[ii] = 1                                  # 平空仓
                idx_ex = np.append(idx_ex, [ii])
                if logmode:
                    print('平空仓, 平仓价%s, longstop:%s, shortstop:%s' % (prc_ex[ii], longstop, shortstop))
            # 如果前一根k线持多仓并未发出平仓信号，则继续持有
            if (sgn_en[ii-1]==0) and (sgn_hd[ii-1] != 0) and (sgn_ex[ii-1] == 0):
                highAfterEntry = max([highAfterEntry, phigh[ii-1]])
                lowAfterEntry = min([lowAfterEntry, plow[ii-1]])
                # 如果前一根K线持仓，当前出现平仓信号
                if (sgn_coverbuy[ii] and (sgn_hd[ii-1]==1)) or (sgn_coversell[ii] and (sgn_hd[ii-1]==-1)):
                    sgn_hd[ii] = sgn_hd[ii-1]
                    prc_en[ii] = prc_en[ii-1]
                    prc_ex[ii] = np.sign(sgn_hd[ii]) * pclose[ii]  # 空头信号， 下一根K线平仓， 平仓价位当前K线收盘价
                    sgn_ex[ii] = -sgn_hd[ii-1]
                    idx_ex = np.append(idx_ex, [ii])
                    if logmode:
                        print('平仓,方向%s, 平仓价%s, longstop:%s, shortstop:%s' % (sgn_ex[ii], prc_ex[ii], longstop, shortstop))
                # 或者如果当前持多仓，最低价低于止损价，以止损价平仓
                elif stoplos and (sgn_hd[ii-1] == 1) and (plow[ii] <= longstop):
                    sgn_hd[ii] = sgn_hd[ii-1]
                    prc_en[ii] = prc_en[ii-1]
                    prc_ex[ii] = np.sign(sgn_hd[ii]) * longstop
                    sgn_ex[ii] = -sgn_hd[ii-1]
                    idx_ex = np.append(idx_ex, [ii])
                    if logmode:
                        print('止损,方向%s, 止损价%s, longstop:%s, shortstop:%s' % (sgn_ex[ii], prc_ex[ii], longstop, shortstop))
                # 或者如果当前持空仓，最高价高于止损价，以止损价平仓
                elif stoplos and (sgn_hd[ii-1] == -1) and (phigh[ii] >= shortstop):
                    sgn_hd[ii] = sgn_hd[ii-1]
                    prc_en[ii] = prc_en[ii-1]
                    prc_ex[ii] = np.sign(sgn_hd[ii]) * shortstop
                    sgn_ex[ii] = -sgn_hd[ii-1]
                    idx_ex = np.append(idx_ex, [ii])
                    if logmode:
                        print('止损,方向%s, 止损价%s, longstop:%s, shortstop:%s' % (sgn_ex[ii], prc_ex[ii], longstop, shortstop))     
                # 无平仓信号，继续持有
                else:
                    sgn_hd[ii] = sgn_hd[ii-1]
                    prc_en[ii] = prc_en[ii-1]
                    prc_ex[ii] = np.sign(sgn_hd[ii]) * pclose[ii]
                    #如果收盘价高于开仓以来最高价，更新多头止损位
                    if pclose[ii] > highAfterEntry:
#                        longstop = min([plow[ii], plow[ii-1]]) - stoplospoint   
                        longstop = min([plow[ii], plow[ii-1]]) - stoplospoint[ii-1]
                    #如果收盘价低于开仓以来最低价，更新空头止损位
                    if pclose[ii] < lowAfterEntry:
#                        shortstop = max([phigh[ii], phigh[ii-1]]) + stoplospoint
                        shortstop = max([phigh[ii], phigh[ii-1]]) + stoplospoint[ii-1]
                    if logmode:
                        print('持仓,方向%s, 最新价%s, longstop:%s, shortstop:%s' % (sgn_hd[ii], prc_ex[ii], longstop, shortstop))
                        
        if (sgn_hd[ii] == 0) or ((sgn_ex[ii] != 0) and (sgn_hd[ii] != 0)):
            if buymode and sgn_buy[ii]:
                sgn_en[ii] = 1
                if logmode:
                    print('触发多头信号')
            if sellmode and sgn_sell[ii]:
                sgn_en[ii] = -1
                if logmode:
                    print('触发空头信号')
    #%%

    if len(idx_en)>len(idx_ex):
        idx_ex=np.append(idx_ex,[len(mktdata)-1])
    #%% 建立交易记录本
    d = {'sgn_en':sgn_en, 'sgn_ex': sgn_ex, 'sgn_hd':sgn_hd,'prc_en':prc_en,'prc_ex':prc_ex}
    tdbook=pd.DataFrame(d, columns=['sgn_en','sgn_ex','sgn_hd','prc_en','prc_ex'])
    tdbook.insert(0,column='datetime',value=mktdata['datetime'])
    if 'symbol' in mktdata.columns:
        tdbook.insert(0,column='symbol',value=mktdata['symbol'])
    #%%
    return tdbook, idx_en, idx_ex
